import pandas as pd
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from rank_bm25 import BM25Okapi
import numpy as np
import re


# 预处理文本：去除停用词并分词
def preprocess_text(text, use_stopwords=True):
    text = re.sub(r'[,.:：&？?—]', '', text.lower())
    if use_stopwords:
        stop_words = stopwords.words('english')
        tokens = word_tokenize(text, 'english')
        return [t for t in tokens if t not in stop_words]
    else:
        tokens = re.sub(r'(?<!\w)-', '', text)
        return tokens.split()


# 计算BM25分数
def calculate_bm25(corpus, query):
    tokenized_corpus = [preprocess_text(doc) for doc in corpus]
    bm25 = BM25Okapi(tokenized_corpus)
    query_tokens = preprocess_text(query)
    return bm25.get_scores(query_tokens)


# 计算Jaccard相似性分数。Jaccard 距离越大，样本相似度越低
def calculate_jaccard(corpus, query):
    query_tokens = set(preprocess_text(query))
    scores = []
    for doc in corpus:
        doc_tokens = set(preprocess_text(doc))
        intersection = len(query_tokens.intersection(doc_tokens))
        union = len(query_tokens.union(doc_tokens))
        j_coef = intersection / union if union != 0 else 0
        score = 1 - j_coef
        scores.append(score)
    return scores


# Reciprocal Rank Fusion (RRF) 方法
def rrf_fusion(ranking1, ranking2, k=60):
    rrf_scores = {}
    for i, doc_id in enumerate(ranking1):
        rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (k + i + 1)
    for i, doc_id in enumerate(ranking2):
        rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (k + i + 1)
    return sorted(rrf_scores, key=rrf_scores.get, reverse=True)


def main(query):
    df = pd.read_excel(r'数据集.xlsx')
    titles = df['title'].tolist()

    # 计算BM25和Jaccard分数
    bm25_scores = calculate_bm25(titles, query)
    jaccard_scores = calculate_jaccard(titles, query)

    # 获取Top50结果(索引)
    # bm25_scores 从大到小 前50
    bm25_top50 = np.argsort(bm25_scores)[::-1][:50]
    # jaccard_scores 从小到大 前50
    jaccard_top50 = np.argsort(jaccard_scores)[:50]

    # RRF融合重排top25(索引)
    top25_rrf = rrf_fusion(bm25_top50, jaccard_top50)[:25]

    # 输出结果
    result = df.iloc[top25_rrf]
    print("RRF重排后的Top25结果：")
    print(result['title'])


if __name__ == "__main__":
    query = "The development of a new generation of information technology"
    main(query)
